#include "model_trans_seg_v2.h"
#include "../public/file_utility.hpp"
#include "../public/dlm_crypt.hpp"
#include "../public/ennx_crypt.hpp"
#include "../txr_algo_dlm_seg_v2/common/common.hpp"
#include "yaml-cpp/yaml.h"

ModelTransSegV2::ModelTransSegV2()
{

}
ModelTransSegV2::~ModelTransSegV2()
{

}

bool ModelTransSegV2::TransModel(st_trans_model_input input)
{
	YAML::Node root = YAML::LoadFile(input.cfg_path);
	YAML::Node config = root["DB"];
	
	std::string onnx_file = config["onnx_file"].as<std::string>();
	std::string engine_file = config["engine_file"].as<std::string>();
	int batch_size = config["BATCH_SIZE"].as<int>();
	int input_channel = config["INPUT_CHANNEL"].as<int>();
	int Max_len = config["Max_len"].as<int>();
	std::vector<float> img_mean = config["img_mean"].as<std::vector<float>>();
	std::vector<float> img_std = config["img_std"].as<std::vector<float>>();
	if (img_mean.size()!=3 || img_std.size()!=3)
	{
		std::cout << "img_mean or img_std cfg error" << std::endl;
	}
	bool dynamic = config["Dynamic"].as<bool>();
	bool dilation = config["Dilation"].as<bool>();
	float Box_thresh = config["Box_thresh"].as<float>();
	float unclip_ratio = config["Unclip_ratio"].as<float>();
	std::string Mode = config["Mode"].as<std::string>();

	std::vector<char> v_model, v_engine;

	namespace fs = boost::filesystem;
	fs::path p(onnx_file.c_str());
	if (p.extension().string() == ".onnx")
	{
		v_model = get_file_data(onnx_file);
	}
	else if (p.extension().string() == ".ennx")
	{
		std::vector<char> v_e_model;
		v_e_model = get_file_data(onnx_file);
		std::string key = "irkdh_573%3?iq5h";
		std::string iv = "1234567890123456";
		v_model = decryptEnnxFile(v_e_model, key, iv);
	}
	else
	{
		return false;
	}
	{
		nvinfer1::ICudaEngine* engine = nullptr;
		onnxToTRTModel(v_model, v_engine, engine, batch_size, Max_len, Max_len, dynamic);
		assert(engine != nullptr);
		if (engine == nullptr)
		{
			return false;
		}
		engine->destroy();
	}

	st_encrypt_info info;
	sprintf_s(info.gpu_model, _countof(info.gpu_model), "%s", input.gpu_model.c_str());
	info.batch_size = batch_size;
	info.input_channel = input_channel;
	info.img_max_len = Max_len;
	for (int i = 0;i<3;++i)
	{
		info.img_mean[i] = img_mean[i];
		info.img_std[i] = img_std[i];
	}
	info.dynamic = dynamic;
	info.dilation = dilation;
	info.box_threshold = Box_thresh;
	info.unclip_ratio = unclip_ratio;
	sprintf_s(info.mode, _countof(info.mode), "%s", Mode.c_str());

	std::string encrypt_file = input.trans_path;
	encryptFile(info, v_engine, encrypt_file);

	return true;
}